// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/table_write_merge.h>

#include "hash_aggregation.h"
#include <pollux/exec/table_writer.h>
#include <pollux/exec/task.h>

namespace kumo::pollux::exec {
    namespace {
        bool isSameCommitContext(
            const melon::Dynamic &first,
            const melon::Dynamic &second) {
            return std::tie(
                       first[TableWriteTraits::kTaskIdContextKey],
                       first[TableWriteTraits::kCommitStrategyContextKey]) ==
                   std::tie(
                       second[TableWriteTraits::kTaskIdContextKey],
                       second[TableWriteTraits::kCommitStrategyContextKey]);
        }

        bool containsNonNullRows(const VectorPtr &vector) {
            if (!vector->may_have_nulls()) {
                return true;
            }
            for (int i = 0; i < vector->size(); ++i) {
                if (!vector->is_null_at(i)) {
                    return true;
                }
            }
            return false;
        }
    } // namespace

    TableWriteMerge::TableWriteMerge(
        int32_t operatorId,
        DriverCtx *driverCtx,
        const std::shared_ptr<const core::TableWriteMergeNode> &tableWriteMergeNode)
        : Operator(
            driverCtx,
            tableWriteMergeNode->outputType(),
            operatorId,
            tableWriteMergeNode->id(),
            "TableWriteMerge") {
        POLLUX_USER_CHECK(outputType_->equivalent(
            *TableWriteTraits::outputType(tableWriteMergeNode->aggregationNode())));
        if (tableWriteMergeNode->aggregationNode() != nullptr) {
            aggregation_ = std::make_unique<HashAggregation>(
                operatorId, driverCtx, tableWriteMergeNode->aggregationNode());
        }
    }

    void TableWriteMerge::initialize() {
        Operator::initialize();
        if (aggregation_ != nullptr) {
            aggregation_->initialize();
        }
    }

    void TableWriteMerge::addInput(RowVectorPtr input) {
        POLLUX_CHECK(!noMoreInput_);
        POLLUX_CHECK_GT(input->size(), 0);

        if (isStatistics(input)) {
            POLLUX_CHECK_NOT_NULL(aggregation_);
            aggregation_->addInput(input);
            return;
        }

        // Increments row count.
        numRows_ += TableWriteTraits::getRowCount(input);

        // Makes sure the lifespan is the same.
        auto commitContext = TableWriteTraits::getTableCommitContext(input);
        if (lastCommitContext_ != nullptr) {
            POLLUX_CHECK(
                isSameCommitContext(lastCommitContext_, commitContext),
                "incompatible table commit context: {} is not compatible with {}",
                lastCommitContext_.asString(),
                commitContext.asString());
        }
        lastCommitContext_ = commitContext;

        // Adds fragments to the buffer. Fragments will be emitted as soon as possible
        // to avoid using extra memory.
        auto fragmentVector = input->childAt(TableWriteTraits::kFragmentChannel);
        if (containsNonNullRows(fragmentVector)) {
            fragmentVectors_.push(fragmentVector);
        }
    }

    void TableWriteMerge::noMoreInput() {
        Operator::noMoreInput();
        if (aggregation_ != nullptr) {
            aggregation_->noMoreInput();
        }
        close();
    }

    RowVectorPtr TableWriteMerge::getOutput() {
        // Passes through fragment pages first to avoid using extra memory.
        if (!fragmentVectors_.empty()) {
            return createFragmentsOutput();
        }

        if (!noMoreInput_ || finished_) {
            return nullptr;
        }

        if (aggregation_ != nullptr && !aggregation_->isFinished()) {
            const std::string commitContext = createTableCommitContext(false);
            return TableWriteTraits::createAggregationStatsOutput(
                outputType_,
                aggregation_->getOutput(),
                StringView(commitContext),
                pool());
        }
        finished_ = true;
        return createLastOutput();
    }

    RowVectorPtr TableWriteMerge::createFragmentsOutput() {
        POLLUX_CHECK(!fragmentVectors_.empty());

        auto outputFragmentVector = fragmentVectors_.front();
        fragmentVectors_.pop();
        const int numOutputRows = outputFragmentVector->size();
        std::vector<VectorPtr> outputColumns(outputType_->size());
        for (int outputChannel = 0; outputChannel < outputType_->size();
             ++outputChannel) {
            if (outputChannel == TableWriteTraits::kFragmentChannel) {
                outputColumns[outputChannel] = std::move(outputFragmentVector);
            } else if (outputChannel == TableWriteTraits::kContextChannel) {
                const std::string commitContext = createTableCommitContext(false);
                outputColumns[outputChannel] =
                        std::make_shared<ConstantVector<StringView> >(
                            pool(),
                            numOutputRows,
                            false /*isNull*/,
                            outputType_->childAt(outputChannel),
                            StringView(commitContext));
            } else {
                outputColumns[outputChannel] = BaseVector::create_null_constant(
                    outputType_->childAt(outputChannel), numOutputRows, pool());
            }
        }
        return std::make_shared<RowVector>(
            pool(), outputType_, nullptr, numOutputRows, outputColumns);
    }

    std::string TableWriteMerge::createTableCommitContext(bool lastOutput) const {
        melon::Dynamic commitContext = lastCommitContext_;
        commitContext[TableWriteTraits::klastPageContextKey] = lastOutput;
        return melon::toJson(commitContext);
    }

    RowVectorPtr TableWriteMerge::createLastOutput() {
        POLLUX_CHECK(
            lastCommitContext_[TableWriteTraits::klastPageContextKey].asBool(),
            "unexpected last table commit context: {}",
            lastCommitContext_.asString());

        auto output = BaseVector::create<RowVector>(outputType_, 1, pool());
        output->resize(1);
        for (int outputChannel = 0; outputChannel < outputType_->size();
             ++outputChannel) {
            if (outputChannel == TableWriteTraits::kRowCountChannel) {
                auto *rowCounterVector =
                        output->childAt(outputChannel)->as_flat_vector<int64_t>();
                rowCounterVector->resize(1);
                rowCounterVector->set(0, numRows_);
            } else if (outputChannel == TableWriteTraits::kContextChannel) {
                auto *contextVector =
                        output->childAt(outputChannel)->as_flat_vector<StringView>();
                contextVector->resize(1);
                const std::string lastCommitContext = createTableCommitContext(true);
                contextVector->set(0, StringView(lastCommitContext));
            } else {
                // All the fragments and statistics shall have already been outputted.
                POLLUX_CHECK(fragmentVectors_.empty());
                output->childAt(outputChannel) = BaseVector::create_null_constant(
                    outputType_->childAt(outputChannel), 1, pool());
            }
        }
        return output;
    }

    bool TableWriteMerge::isStatistics(RowVectorPtr input) {
        return input->childAt(TableWriteTraits::kRowCountChannel)->is_null_at(0) &&
               input->childAt(TableWriteTraits::kFragmentChannel)->is_null_at(0);
    }
} // namespace kumo::pollux::exec
